from collections import OrderedDict
from functools import partial
from typing import Callable, Optional

import torch.nn as nn
import torch
from torch import Tensor
import numpy
import torchaudio
import torch.nn.functional as F

from models.unit import Conv2d, Reduction_A


class Stem(nn.Module):
    def __init__(self, in_channels):
        super(Stem, self).__init__()
        self.features = nn.Sequential(
            Conv2d(in_channels, 32, 3, stride=2, padding=0, bias=False),  # 149 x 149 x 32
            Conv2d(32, 32, 3, stride=1, padding=0, bias=False),  # 147 x 147 x 32
            Conv2d(32, 64, 3, stride=1, padding=1, bias=False),  # 147 x 147 x 64
            nn.MaxPool2d(3, stride=2, padding=0),  # 73 x 73 x 64
            Conv2d(64, 80, 1, stride=1, padding=0, bias=False),  # 73 x 73 x 80
            Conv2d(80, 192, 3, stride=1, padding=0, bias=False),  # 71 x 71 x 192
            nn.MaxPool2d(3, stride=2, padding=0),  # 35 x 35 x 192
        )
        self.branch_0 = Conv2d(192, 96, 1, stride=1, padding=0, bias=False)
        self.branch_1 = nn.Sequential(
            Conv2d(192, 48, 1, stride=1, padding=0, bias=False),
            Conv2d(48, 64, 5, stride=1, padding=2, bias=False),
        )
        self.branch_2 = nn.Sequential(
            Conv2d(192, 64, 1, stride=1, padding=0, bias=False),
            Conv2d(64, 96, 3, stride=1, padding=1, bias=False),
            Conv2d(96, 96, 3, stride=1, padding=1, bias=False),
        )
        self.branch_3 = nn.Sequential(
            nn.AvgPool2d(3, stride=1, padding=1, count_include_pad=False),
            Conv2d(192, 64, 1, stride=1, padding=0, bias=False)
        )

    def forward(self, x):
        x = self.features(x)
        x0 = self.branch_0(x)
        x1 = self.branch_1(x)
        x2 = self.branch_2(x)
        x3 = self.branch_3(x)
        return torch.cat((x0, x1, x2, x3), dim=1)


class Inception_ResNet_A(nn.Module):
    def __init__(self, in_channels, scale=1.0):
        super(Inception_ResNet_A, self).__init__()
        self.scale = scale
        self.branch_0 = Conv2d(in_channels, 32, 1, stride=1, padding=0, bias=False)
        self.branch_1 = nn.Sequential(
            Conv2d(in_channels, 32, 1, stride=1, padding=0, bias=False),
            Conv2d(32, 32, 3, stride=1, padding=1, bias=False)
        )
        self.branch_2 = nn.Sequential(
            Conv2d(in_channels, 32, 1, stride=1, padding=0, bias=False),
            Conv2d(32, 48, 3, stride=1, padding=1, bias=False),
            Conv2d(48, 64, 3, stride=1, padding=1, bias=False)
        )
        self.conv = nn.Conv2d(128, 320, 1, stride=1, padding=0, bias=True)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        x0 = self.branch_0(x)
        x1 = self.branch_1(x)
        x2 = self.branch_2(x)
        x_res = torch.cat((x0, x1, x2), dim=1)
        x_res = self.conv(x_res)
        return self.relu(x + self.scale * x_res)


class Inception_ResNet_B(nn.Module):
    def __init__(self, in_channels, scale=1.0):
        super(Inception_ResNet_B, self).__init__()
        self.scale = scale
        self.branch_0 = Conv2d(in_channels, 192, 1, stride=1, padding=0, bias=False)
        self.branch_1 = nn.Sequential(
            Conv2d(in_channels, 128, 1, stride=1, padding=0, bias=False),
            Conv2d(128, 160, (1, 7), stride=1, padding=(0, 3), bias=False),
            Conv2d(160, 192, (7, 1), stride=1, padding=(3, 0), bias=False)
        )
        self.conv = nn.Conv2d(384, 1088, 1, stride=1, padding=0, bias=True)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        x0 = self.branch_0(x)
        x1 = self.branch_1(x)
        x_res = torch.cat((x0, x1), dim=1)
        x_res = self.conv(x_res)
        return self.relu(x + self.scale * x_res)


class Reduciton_B(nn.Module):
    def __init__(self, in_channels):
        super(Reduciton_B, self).__init__()
        self.branch_0 = nn.Sequential(
            Conv2d(in_channels, 256, 1, stride=1, padding=0, bias=False),
            Conv2d(256, 384, 3, stride=2, padding=0, bias=False)
        )
        self.branch_1 = nn.Sequential(
            Conv2d(in_channels, 256, 1, stride=1, padding=0, bias=False),
            Conv2d(256, 288, 3, stride=2, padding=0, bias=False),
        )
        self.branch_2 = nn.Sequential(
            Conv2d(in_channels, 256, 1, stride=1, padding=0, bias=False),
            Conv2d(256, 288, 3, stride=1, padding=1, bias=False),
            Conv2d(288, 320, 3, stride=2, padding=0, bias=False)
        )
        self.branch_3 = nn.MaxPool2d(3, stride=2, padding=0)

    def forward(self, x):
        x0 = self.branch_0(x)
        x1 = self.branch_1(x)
        x2 = self.branch_2(x)
        x3 = self.branch_3(x)
        return torch.cat((x0, x1, x2, x3), dim=1)


class Inception_ResNet_C(nn.Module):
    def __init__(self, in_channels, scale=1.0, activation=True):
        super(Inception_ResNet_C, self).__init__()
        self.scale = scale
        self.activation = activation
        self.branch_0 = Conv2d(in_channels, 192, 1, stride=1, padding=0, bias=False)
        self.branch_1 = nn.Sequential(
            Conv2d(in_channels, 192, 1, stride=1, padding=0, bias=False),
            Conv2d(192, 224, (1, 3), stride=1, padding=(0, 1), bias=False),
            Conv2d(224, 256, (3, 1), stride=1, padding=(1, 0), bias=False)
        )
        self.conv = nn.Conv2d(448, 2080, 1, stride=1, padding=0, bias=True)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        x0 = self.branch_0(x)
        x1 = self.branch_1(x)
        x_res = torch.cat((x0, x1), dim=1)
        x_res = self.conv(x_res)
        if self.activation:
            return self.relu(x + self.scale * x_res)
        return x + self.scale * x_res


class Inception_ResNetv2(nn.Module):
    def __init__(self,
                 in_channels=3,
                 classes=1000,
                 k=256,
                 l=256,
                 m=384,
                 n=384,
                 num_filters=[32, 64, 128, 256],
                 nOut=512,
                 encoder_type='SAP',
                 n_mels=40,
                 log_input=True, **kwargs):
        super(Inception_ResNetv2, self).__init__()

        print('Embedding size is %d, encoder %s.' % (nOut, encoder_type))

        # self.inplanes = num_filters[0]
        self.encoder_type = encoder_type
        self.n_mels = n_mels
        self.log_input = log_input
        # self.conv1 = nn.Conv2d(1, num_filters[0], kernel_size=7, stride=(2, 1), padding=3,
        #                        bias=False)
        # self.bn1 = nn.BatchNorm2d(num_filters[0])
        # self.relu = nn.ReLU(inplace=True)
        self.instancenorm = nn.InstanceNorm1d(n_mels)
        self.torchfb = torchaudio.transforms.MelSpectrogram(sample_rate=16000, n_fft=512, win_length=400,
                                                            hop_length=160, window_fn=torch.hamming_window,
                                                            n_mels=n_mels)
        if self.encoder_type == "SAP":
            self.sap_linear = nn.Linear(640, 160)
            self.attention = self.new_parameter(160, 1)
            out_dim = 640
        elif self.encoder_type == "ASP":
            self.sap_linear = nn.Linear(num_filters[3] * 4, num_filters[3] * 4)
            self.attention = self.new_parameter(num_filters[3] * 4, 1)
            out_dim = num_filters[3] * 4 * 2
        else:
            raise ValueError('Undefined encoder')

        self.fc = nn.Linear(1280, nOut)

        blocks = []
        blocks.append(Stem(in_channels))
        for i in range(10):
            blocks.append(Inception_ResNet_A(320, 0.17))
        blocks.append(Reduction_A(320, k, l, m, n))
        for i in range(20):
            blocks.append(Inception_ResNet_B(1088, 0.10))
        blocks.append(Reduciton_B(1088))
        for i in range(9):
            blocks.append(Inception_ResNet_C(2080, 0.20))
        blocks.append(Inception_ResNet_C(2080, activation=False))
        self.features = nn.Sequential(*blocks)
        self.conv = Conv2d(2080, 1536, 1, stride=1, padding=0, bias=False)
        self.global_average_pooling = nn.AdaptiveAvgPool2d((1, 1))
        self.linear = nn.Linear(1536, classes)

    def new_parameter(self, *size):
        out = nn.Parameter(torch.FloatTensor(*size))
        nn.init.xavier_normal_(out)
        return out


def forward(self, x):
    # with torch.no_grad():
    #     with torch.cuda.amp.autocast(enabled=False):
    #         x = self.torchfb(x) + 1e-6
    #         if self.log_input: x = x.log()
    #         x = self.instancenorm(x).unsqueeze(1).detach()

    x = self.features(x)
    x = self.conv(x)
    x = self.global_average_pooling(x)
    x = x.view(x.size(0), -1)
    x = self.linear(x)
    return x


class EfficientNetV2_mel(nn.Module):
    def __init__(self,
                 model_cnf: list,
                 dropout_rate: float = 0.2,
                 drop_connect_rate: float = 0.2,
                 num_features: int = 1280,
                 num_filters=[32, 64, 128, 256],
                 nOut=512,
                 encoder_type='SAP',
                 n_mels=40,
                 log_input=True, **kwargs):
        super(EfficientNetV2_mel, self).__init__()

        print('Embedding size is %d, encoder %s.' % (nOut, encoder_type))

        self.inplanes = num_filters[0]
        self.encoder_type = encoder_type
        self.n_mels = n_mels
        self.log_input = log_input

        self.conv1 = nn.Conv2d(1, num_filters[0], kernel_size=7, stride=(2, 1), padding=3,
                               bias=False)
        self.bn1 = nn.BatchNorm2d(num_filters[0])
        self.relu = nn.ReLU(inplace=True)

        self.instancenorm = nn.InstanceNorm1d(n_mels)
        self.torchfb = torchaudio.transforms.MelSpectrogram(sample_rate=16000, n_fft=512, win_length=400,
                                                            hop_length=160, window_fn=torch.hamming_window,
                                                            n_mels=n_mels)

        if self.encoder_type == "SAP":
            self.sap_linear = nn.Linear(640, 160)
            self.attention = self.new_parameter(160, 1)
            out_dim = 640
        elif self.encoder_type == "ASP":
            self.sap_linear = nn.Linear(num_filters[3] * 4, num_filters[3] * 4)
            self.attention = self.new_parameter(num_filters[3] * 4, 1)
            out_dim = num_filters[3] * 4 * 2
        else:
            raise ValueError('Undefined encoder')

        self.fc = nn.Linear(1280, nOut)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

        # 以下是effitientNet 内容
        for cnf in model_cnf:
            assert len(cnf) == 8

        norm_layer = partial(nn.BatchNorm2d, eps=1e-3, momentum=0.1)

        stem_filter_num = model_cnf[0][4]

        # 更改 3-》1
        self.stem = ConvBNAct(1,
                              stem_filter_num,
                              kernel_size=3,
                              stride=2,
                              norm_layer=norm_layer)  # 激活函数默认是SiLU

        total_blocks = sum([i[0] for i in model_cnf])
        block_id = 0
        blocks = []
        for cnf in model_cnf:
            repeats = cnf[0]
            op = FusedMBConv if cnf[-2] == 0 else MBConv
            for i in range(repeats):
                blocks.append(op(kernel_size=cnf[1],
                                 input_c=cnf[4] if i == 0 else cnf[5],
                                 out_c=cnf[5],
                                 expand_ratio=cnf[3],
                                 stride=cnf[2] if i == 0 else 1,
                                 se_ratio=cnf[-1],
                                 drop_rate=drop_connect_rate * block_id / total_blocks,
                                 norm_layer=norm_layer))
                block_id += 1
        self.blocks = nn.Sequential(*blocks)

        head_input_c = model_cnf[-1][-3]
        head = OrderedDict()

        head.update({"project_conv": ConvBNAct(head_input_c,
                                               num_features,
                                               kernel_size=1,
                                               norm_layer=norm_layer)})  # 激活函数默认是SiLU

        head.update({"avgpool": nn.AdaptiveAvgPool2d(1)})
        head.update({"flatten": nn.Flatten()})

        if dropout_rate > 0:
            head.update({"dropout": nn.Dropout(p=dropout_rate, inplace=True)})
        # head.update({"classifier": nn.Linear(num_features, num_classes)})

        self.head = nn.Sequential(head)

        # initial weights
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode="fan_out")
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.zeros_(m.bias)

    def new_parameter(self, *size):
        out = nn.Parameter(torch.FloatTensor(*size))
        nn.init.xavier_normal_(out)
        return out

    def forward(self, x):

        with torch.no_grad():
            with torch.cuda.amp.autocast(enabled=False):
                x = self.torchfb(x) + 1e-6
                if self.log_input: x = x.log()
                x = self.instancenorm(x).unsqueeze(1).detach()

        x = self.stem(x)
        x = self.blocks(x)
        # self.blocks(x)(x): torch.Size([200, 640, 2, 7])
        # print('self.blocks(x)(x):', numpy.shape(x))
        x = self.head(x)
        # self.head(x+3x): torch.Size([200, 1280])
        # print('self.head(x):', x.shape)
        # exit()
        # x = torch.mean(x, dim=2, keepdim=True)
        # torch.mean(x, dim=2, keepdim=True): torch.Size([200, 640, 1, 7])
        # print('torch.mean(x, dim=2, keepdim=True):', numpy.shape(x))

        # if self.encoder_type == "SAP":
        #     x = x.permute(0, 3, 1, 2).squeeze(-1)
        #     # x.permute(0, 3, 1, 2).squeeze(-1) torch.Size([200, 7, 640])
        #     # print('x.permute(0, 3, 1, 2).squeeze(-1)', numpy.shape(x))
        #
        #     h = torch.tanh(self.sap_linear(x))
        #     w = torch.matmul(h, self.attention).squeeze(dim=2)
        #     w = F.softmax(w, dim=1).view(x.size(0), x.size(1), 1)
        #     x = torch.sum(x * w, dim=1)
        # elif self.encoder_type == "ASP":
        #     x = x.permute(0, 3, 1, 2).squeeze(-1)
        #     h = torch.tanh(self.sap_linear(x))
        #     w = torch.matmul(h, self.attention).squeeze(dim=2)
        #     w = F.softmax(w, dim=1).view(x.size(0), x.size(1), 1)
        #     mu = torch.sum(x * w, dim=1)
        #     rh = torch.sqrt((torch.sum((x ** 2) * w, dim=1) - mu ** 2).clamp(min=1e-5))
        #     x = torch.cat((mu, rh), 1)
        # print('SAP:', numpy.shape(x))
        # exit()
        x = x.view(x.size()[0], -1)
        # x.view(x.size()[0], -1): torch.Size([200, 1280])
        # print('x.view(x.size()[0], -1):', numpy.shape(x))
        x = self.fc(x)
        # self.fc(x): torch.Size([200, 512])
        # print('self.fc(x):', numpy.shape(x))
        # exit()
        return x


def MainModel1(nOut=256, **kwargs):
    # Number of filters
    model_config = [[2, 3, 1, 1, 24, 24, 0, 0],
                    [4, 3, 2, 4, 24, 48, 0, 0],
                    [4, 3, 2, 4, 48, 64, 0, 0],
                    [6, 3, 2, 4, 64, 128, 1, 0.25],
                    [9, 3, 1, 6, 128, 160, 1, 0.25],
                    [15, 3, 2, 6, 160, 256, 1, 0.25]]
    num_filters = [32, 64, 128, 256]
    model = EfficientNetV2_mel(model_cnf=model_config,
                               nOut=nOut,
                               num_filters=num_filters,
                               dropout_rate=0.4,
                               **kwargs)

    return model


def MainModel(nOut=256, **kwargs):
    # Number of filters

    model_config = [[2, 3, 1, 1, 24, 24, 0, 0],
                    [4, 3, 2, 4, 24, 48, 0, 0],
                    [4, 3, 2, 4, 48, 64, 0, 0],
                    [6, 3, 2, 4, 64, 128, 1, 0.25],
                    [9, 3, 1, 6, 128, 160, 1, 0.25],
                    [15, 3, 2, 6, 160, 256, 1, 0.25]]
    num_filters = [32, 64, 128, 256]
    model = Inception_ResNetv2()
    # model = EfficientNetV2_mel(model_cnf=model_config,
    #                            nOut=nOut,
    #                            num_filters=num_filters,
    #                            dropout_rate=0.4,
    #                            **kwargs)
    # print('inti OK!!!')
    # exit()
    return model


if __name__ == '__main__':
    model = MainModel().cuda()

    x = torch.rand(200, 1, 20, 202).cuda()
    r = model(x)
